import json
import os


def convert_coordinates(coord):
    """Convert [x1,y1,x2,y2] format to COCO format [x,y,width,height]"""
    x1, y1 = coord[0]
    x2, y2 = coord[1]
    width = x2 - x1
    height = y2 - y1
    return [x1, y1, width, height]


def parse_line(line):
    """Parse a single line of the input format"""
    parts = line.strip().split('\t')
    image_name = parts[0]
    annotations = []

    for part in parts[1:]:
        try:
            data = json.loads(part)
            annotations.append({
                'category': data['value'],
                'bbox': convert_coordinates(data['coordinate'])
            })
        except json.JSONDecodeError:
            print(f"Warning: Could not parse annotation: {part}")
            continue

    return image_name, annotations


def convert_to_coco(input_file, image_dir):
    """Convert the dataset to COCO format"""
    # Initialize COCO format structure
    dataset = {
        'images': [],
        'annotations': [],
        'categories': []
    }

    # Create category mapping
    categories = {'bolt': 1, 'nut': 2}
    dataset['categories'] = [
        {'id': 1, 'name': 'bolt'},
        {'id': 2, 'name': 'nut'}
    ]

    annotation_id = 1
    with open(input_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            image_name, annotations = parse_line(line)

            # Add image info
            image_id = line_num
            # You might want to get actual image dimensions here
            dataset['images'].append({
                'id': image_id,
                'file_name': image_name,
                'width': 1440,  # Replace with actual image width
                'height': 1080  # Replace with actual image height
            })

            # Add annotations
            for ann in annotations:
                dataset['annotations'].append({
                    'id': annotation_id,
                    'image_id': image_id,
                    'category_id': categories[ann['category']],
                    'bbox': ann['bbox'],
                    'area': ann['bbox'][2] * ann['bbox'][3],
                    'iscrowd': 0
                })
                annotation_id += 1

    return dataset


def convert_to_voc(input_file, output_dir):
    """Convert the dataset to Pascal VOC format"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    def create_xml(image_name, annotations, width=1600, height=1200):
        xml_content = f"""<?xml version="1.0" encoding="UTF-8"?>
<annotation>
    <filename>{image_name}</filename>
    <size>
        <width>{width}</width>
        <height>{height}</height>
        <depth>3</depth>
    </size>
"""
        for ann in annotations:
            bbox = ann['bbox']
            xml_content += f"""    <object>
        <name>{ann['category']}</name>
        <bndbox>
            <xmin>{int(bbox[0])}</xmin>
            <ymin>{int(bbox[1])}</ymin>
            <xmax>{int(bbox[0] + bbox[2])}</xmax>
            <ymax>{int(bbox[1] + bbox[3])}</ymax>
        </bndbox>
    </object>
"""
        xml_content += "</annotation>"
        return xml_content

    with open(input_file, 'r', encoding='utf-8') as f:
        for line in f:
            image_name, annotations = parse_line(line)
            xml_content = create_xml(image_name, annotations)

            # Save XML file
            xml_filename = os.path.splitext(image_name)[0] + '.xml'
            xml_path = os.path.join(output_dir, xml_filename)
            with open(xml_path, 'w', encoding='utf-8') as xml_file:
                xml_file.write(xml_content)


# Usage example
if __name__ == "__main__":
    input_file = "/Users/jiangfeng/PycharmProjects/Net/Faster_RCNN/lslm/lslm/train.txt"

    # Convert to COCO format
    coco_dataset = convert_to_coco(input_file, "images/")
    with open("annotations_coco.json", 'w') as f:
        json.dump(coco_dataset, f, indent=2)

    # Convert to VOC format
    convert_to_voc(input_file, "annotations_voc/")
